/**
 * @file stress.c
 * @author snow (zziywang@163.com)
 * @brief { life is too short to learn cpp! }
 * @version 1.0
 * @date 2022-05-23
 *
 * @copyright Copyright (c) 2022
 *
 */

#include <arpa/inet.h>
#include <assert.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/epoll.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

static const char *request = "xxxxxxxxxxxx\n";

// static const char *request = "GET http://localhost/index.html HTTP/1.1\r\nConnection: keep-alive\r\n\r\nxxxxxxxxxxxx";

int setNonBlocking(int fd)
{
    int oldOption = fcntl(fd, F_GETFL);
    int newOption = oldOption | O_NONBLOCK;
    fcntl(fd, F_SETFL, newOption);
    return oldOption;
}

void addFd(int epollFd, int fd)
{
    struct epoll_event event;
    event.data.fd = fd;
    event.events = EPOLLOUT | EPOLLET | EPOLLERR;
    epoll_ctl(epollFd, EPOLL_CTL_ADD, fd, &event);
    setNonBlocking(fd);
}

// 向服务器写len字节数据
bool writenBytes(int socketFd, const char *buf, int len)
{
    int bytesWrite = 0;
    printf("write out %d bytes to socket %d\n", len, socketFd);
    while (true) {
        bytesWrite = send(socketFd, buf, len, 0);
        if (-1 == bytesWrite) {
            return false;
        } else if (0 == bytesWrite) {
            return false;
        }

        len -= bytesWrite;
        buf = buf + bytesWrite;
        if (len <= 0) {
            return true;
        }
    }
}

bool readOnce(int socketFd, char *buf, int len)
{
    int bytesRead = 0;
    memset(buf, 0, len);
    bytesRead = recv(socketFd, buf, len, 0);
    if (-1 == bytesRead) { // 错误
        return false;
    } else if (0 == bytesRead) { // 主动断开
        return false;
    }
    printf("read in %d bytes from socket %d with content: %s\n", bytesRead, socketFd, buf);
    return true;
}

// 每隔1s创建num个连接
void startConnection(int epollFd, int num, const char *ip, int port)
{
    int ret = 0;
    struct sockaddr_in address;
    memset(&address, 0, sizeof(address));
    address.sin_family = AF_INET;
    inet_pton(AF_INET, ip, &address.sin_addr);
    address.sin_port = htons(port);

    for (int i = 0; i < num; ++i) {
        sleep(1);
        int socketFd = socket(PF_INET, SOCK_STREAM, 0);
        printf("create 1 sock\n");
        if (socketFd < 0) {
            continue;
        }
        if (connect(socketFd, (struct sockaddr *)&address, sizeof(address)) == 0) {
            printf("build connection %d\n", i + 1);
            addFd(epollFd, socketFd);
        }
    }
}

void closeConnection(int epollFd, int socketFd)
{
    epoll_ctl(epollFd, EPOLL_CTL_DEL, socketFd, 0);
    close(socketFd);
}

int main(int argc, char **argv)
{
    // ./stress ip port 连接数
    assert(argc == 4);
    int epollFd = epoll_create1(0);
    startConnection(epollFd, atoi(argv[3]), argv[1], atoi(argv[2]));
    struct epoll_event events[10000];
    char buf[2048];
    // char request[2048] = "This is a stress test";
    while (true) {
        int fds = epoll_wait(epollFd, events, 10000, 2000); // 20s超时
        for (int i = 0; i < fds; ++i) {
            int socketFd = events[i].data.fd;
            if (events[i].events & EPOLLIN) {
                if (!readOnce(socketFd, buf, 2048)) {
                    closeConnection(epollFd, socketFd);
                }
                struct epoll_event event;
                event.events = EPOLLOUT | EPOLLET | EPOLLERR;
                event.data.fd = socketFd;
                epoll_ctl(epollFd, EPOLL_CTL_MOD, socketFd, &event);
            } else if (events[i].events & EPOLLOUT) {
                if (!writenBytes(socketFd, request, strlen(request))) {
                    closeConnection(epollFd, socketFd);
                }
                struct epoll_event event;
                event.events = EPOLLOUT | EPOLLET | EPOLLERR;
                event.data.fd = socketFd;
                epoll_ctl(epollFd, EPOLL_CTL_MOD, socketFd, &event);
            } else if (events[i].events & EPOLLERR) {
                closeConnection(epollFd, socketFd);
            }
        }
    }

    return 0;
}